package xyz.naokeziteng.leetcode;

import java.util.Random;

/**
 * 215: 第K个最大元素
 */
public class FindKthLargest {

    private static Random RANDOM = new Random();

    public static int findKthLargest(int[] nums, int k) {
        return findKthLargest3(nums, 0, nums.length - 1, k);
    }

    private static int findKthLargest2(int[] arr, int l, int r, int k) {
        swap(arr, l, l + RANDOM.nextInt(r - l + 1));
        int v = arr[l];
        int j = l;
        for (int i = l + 1; i <= r; i++) {
            if (arr[i] < v) {
                continue;
            }
            j++;
            swap(arr, i, j);
        }
        swap(arr, l, j);
        if (j > k - 1) {
            return findKthLargest2(arr, l, j - 1, k);
        } else if (j < k - 1) {
            return findKthLargest2(arr, j + 1, r, k);
        } else {
            return arr[j];
        }
    }


    private static int findKthLargest3(int[] arr, int l, int r, int k) {
        while (true) {
            swap(arr, l, l + RANDOM.nextInt(r - l + 1));
            int v = arr[l];
            int j = l;
            for (int i = l + 1; i <= r; i++) {
                if (arr[i] < v) {
                    continue;
                }
                j++;
                swap(arr, i, j);
            }
            swap(arr, l, j);
            if (j == k - 1) {
                return arr[j];
            } else if (j > k - 1) {
                r = j - 1;
            } else {
                l = j + 1;
            }
        }
    }

    private static int findKthLargest(int[] arr, int l, int r, int k) {
        swap(arr, l, l + RANDOM.nextInt(r - l + 1));
        int v = arr[l];
        int lt = r + 1;
        int i = l + 1;
        int gt = l;
        while (i < lt) {
            if (arr[i] < v) {
                lt--;
                swap(arr, i, lt);
            } else if (arr[i] > v) {
                gt++;
                swap(arr, i, gt);
                i++;
            } else {
                i++;
            }
        }
        swap(arr, l, gt);
        if (k - 1 < gt) {
            return findKthLargest(arr, l, gt - 1, k);
        } else if (k - 1 > lt - 1) {
            return findKthLargest(arr, lt, r, k);
        } else {
            int tmp = gt;
            while (tmp < lt) {
                if (k - 1 == tmp) {
                    break;
                }
                tmp++;
            }
            return arr[tmp];
        }
    }

    private static void swap(int[] arr, int i, int j) {
        int tmp = arr[i];
        arr[i] = arr[j];
        arr[j] = tmp;
    }

    public static void main(String[] args) {
        int[] arr = {3, 2, 1, 5, 6, 4};
        System.out.println(findKthLargest(arr, 2));;
    }

}
